package com.chaos.auth.config;


import com.chaos.framework.model.dto.result.data.DataResult;
import com.chaos.framework.model.enums.ResultEnum;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.exceptions.OAuth2Exception;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

@Order(1)
@ControllerAdvice
@Slf4j
public class Oauth2ResponseBodyAdvice implements ResponseBodyAdvice<Object> {
    @Override
    public boolean supports(MethodParameter methodParameter, Class aClass) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter methodParameter, MediaType mediaType, Class aClass, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {

        if (body instanceof DefaultOAuth2AccessToken) {
            DataResult<Object> result = DataResult.ok(body);
            result.setData(body);
            return result;
        } else if (body instanceof OAuth2Exception) {
            // 处理错误响应
            OAuth2Exception oAuth2Exception = (OAuth2Exception) body;
            DataResult<Object> result = DataResult.error(ResultEnum.INVALID_REQUEST.getCode(), oAuth2Exception.getMessage());
            return result;
        }
        return body;

    }


}

